"""Model inference.

Use OM model to infer the original images, and then save the prediction
results or visualization results.

Example:
    >>> python tools/infer.py \
    >>>     mmdetection/configs/yolox/yolox_s_8xb8-300e_coco.py \
    >>>     models/yolox_s_bs1.om \
    >>>     --inputs demo/test_images \
    >>>     --out_dir demo/outputs
"""

from argparse import ArgumentParser
from mmengine.logging import print_log
from npu_mmdet.apis import DetInferencer


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('config', type=str, help='Path to model config file.')
    parser.add_argument('om_model', type=str, help='Path to OM model.')
    parser.add_argument('--inputs', type=str,
                        help='One image file or folder path.')
    parser.add_argument('--out_dir', type=str,
                        default='outputs',
                        help='Output directory to save images or predictions.')
    parser.add_argument('--device', type=int,
                        default=0,
                        help='Npu id used for inference.')
    parser.add_argument('--pred_score_thr', type=float,
                        default=0.3,
                        help='bbox score threshold.')
    parser.add_argument('--show', action='store_true',
                        help='Display the image in a popup window.')
    parser.add_argument('--no_save_vis', action='store_true',
                        help='Do not save detection vis results')
    parser.add_argument('--no_save_pred', action='store_true',
                        help='Do not save detection json results')
    parser.add_argument('--print_result', action='store_true',
                        help='Whether to print the results.')

    call_args = vars(parser.parse_args())

    if call_args['no_save_vis'] and call_args['no_save_pred']:
        call_args['out_dir'] = ''

    init_kws = ['config', 'om_model', 'device']
    init_args = {}
    for init_kw in init_kws:
        init_args[init_kw] = call_args.pop(init_kw)

    return init_args, call_args


def main():
    init_args, call_args = parse_args()
    inferencer = DetInferencer(**init_args)
    inferencer(**call_args)

    if call_args['out_dir'] != '' and \
            not (call_args['no_save_vis'] and call_args['no_save_pred']):
        print_log(f'results have been saved at {call_args["out_dir"]}')


if __name__ == '__main__':
    main()
